from yolov5 import YOLOv5
import cv2
import pdb
import time
import numpy as np

def compute_homography(H_ratio, V_ratio, im_size):
    rationed_hight = im_size[1] * V_ratio
    rationed_width = im_size[0] * H_ratio
    src = np.array([[0, 0], [0, im_size[1]], [
        im_size[0], im_size[1]], [im_size[0], 0]])
    dst = np.array([[0 + rationed_width / 2, 0 + rationed_hight], [0, im_size[1]], [im_size[0],
                                                                                    im_size[1]],
                    [im_size[0] - rationed_width / 2, 0 + rationed_hight]], np.int32)
    h, status = cv2.findHomography(src, dst)
    return h

def ellipse_polyline(ellipse, n=32):
    """Convert ellipse into a polyline, for each ellipse

    Args:
        ellipse (list): ellipse list
        n (int, optional): segments. Defaults to 100.

    Returns:
        list: [description]
    """

    t = np.linspace(0, 2 * np.pi, n, endpoint=False)
    st = np.sin(t)
    ct = np.cos(t)

    x0, y0, a, b = ellipse
    p = np.empty((n, 2))
    p[:, 0] = x0 + a * ct
    p[:, 1] = y0 + b * st

    return p

def intersections(a, b):
    """check if two polylines are intersected

    Args:
        a (polyline): polyline a
        b (polyline): polyline b

    Returns:
        boolean: true if polylines are intersected
    """
    try:
        ea = LinearRing(a)
        eb = LinearRing(b)
        return ea.intersection(eb)
    except:
        return False

def trans_box_to_ell(rect): # 中心,long,short
    left, right, top, bottom = rect
    center_x = round((left+right)/2)
    center_y = round((top+bottom)/2)
    long= right - center_x
    short= bottom -center_y
    return center_x,center_y,long,short

def compute_overlap(rect_1, rect_2,use_ell_box=True):  # rect = [left,right,top,bottom]
    if use_ell_box:
        x_overlap = max(0, min(rect_1[1], rect_2[1]) - max(rect_1[0], rect_2[0]))
        y_overlap = max(0, min(rect_1[3], rect_2[3]) - max(rect_1[2], rect_2[2]))
        overlapArea = x_overlap * y_overlap

    if use_ell_box==False:
        ell1 = trans_box_to_ell(rect_1)
        ell2 = trans_box_to_ell(rect_2)
        p1 = ellipse_polyline(ell1)
        p2 = ellipse_polyline(ell2)
        overlapArea = intersections(p1,p2)

    if overlapArea:
        overlaps = True
    else:
        overlaps = False
    return overlaps

def compute_overlap_gt(pt1, pt2,scale):
    dist = distance.cdist(pt1, pt2)/scale
    if dist <= 2:
        overlaps = True
    else:
        overlaps = False
    return overlaps

def get_ells_from_person_bboxes(im,list_bboxes,horizontal_ratio, vertical_ratio):
    '''
    :param im:
    :param list_bboxes: [(l,t,r,b,cls_name,ID)]
    :param horizontal_ratio:
    :param vertical_ratio:
    :return: draw_ellipse_requirements[(centre, longaxe,shortaxe,ID)]
    '''

    im_size = (im.shape[1], im.shape[0])
    draw_ellipse_requirements=[]
    ellipse_boxes=[]
    center_world_coordinates = []

    for bbox in list_bboxes: #x1, y1, x2, y2, lbl, conf
        left, top,right, bottom ,cls_name,ID= bbox
        bb_center = np.array([(left + right) / 2, (top + bottom) / 2], np.int32)
        calculated_hight = round((bottom - top)/body_height)# 人周围1meter 对应的图像高度   # 像素高/ 人高1.6meter = x 像素 / 1 meter

        homography_matrix = compute_homography(
            horizontal_ratio, vertical_ratio, im_size)

        pts = np.array([[bb_center[0], top], [bb_center[0], bottom]], np.float32)
        pts1 = pts.reshape(-1, 1, 2).astype(np.float32)  # (n, 1, 2)
        dst1 = cv2.perspectiveTransform(pts1, homography_matrix)
        width = int(dst1[1, 0][1] - dst1[0, 0][1])
        # Basic control to avoid very wrong ellipses -> To be fixed
        if width > 0.5 * calculated_hight:
            width = int(0.5 * calculated_hight)
        ellipse_bbx = [bb_center[0] - calculated_hight,
                       bb_center[0] + calculated_hight, bottom - width, bottom + width]
        draw_ellipse_requirements.append(
            [bb_center[0], bottom, calculated_hight, width,ID])  # h = longaxe,  w= shortaxe
        ellipse_boxes.append(ellipse_bbx)
        center_pixel = np.array([[[bb_center[0], bottom]]], np.float32)
        center_world = cv2.perspectiveTransform(center_pixel, H_truth)
        center_world_coordinates.append(center_world[0])
    return draw_ellipse_requirements,ellipse_boxes,center_world_coordinates

def to_rectangle(ellipse):
        """Convert ellipse to rectangle (top, left, bottom, right)

        Args:
            ellipse (tuple): bounding rectangle descriptor
        """
        x, y, a, b ,_= ellipse
        return(x-a,y-b,x+a,y+b)

def contains(ellipse1, ellipse2):
        """check if ellipses bounding rectangles are overlapped

        Args:
            ellipse1 (tuple): ellipse one
            ellipse2 (tuple): ellipse two

        Returns:
            boolean:
        """
        r1 = to_rectangle(ellipse1)
        r2 = to_rectangle(ellipse2)

        if (r1[0]>=r2[2]) or (r1[2]<=r2[0]) or (r1[3]<=r2[1]) or (r1[1]>=r2[3]):
            return False
        else:
            return True

def overlap_infer(im, draw_ellipse_requirements, ellipse_boxes):
    is_skeletal_overlapped = np.zeros(len(ellipse_boxes))

    # Find a body,if is already in violation do not consider it
    for body_idx_1 in range(len(draw_ellipse_requirements)):
        if is_skeletal_overlapped[body_idx_1] == 1:
            continue

        # Find another body (if it already in violation it is the same)
        for body_idx_2 in range(len(draw_ellipse_requirements)):
            # Same body, skip
            if body_idx_1 == body_idx_2:
                continue

            ellipse1 = tuple(draw_ellipse_requirements[body_idx_1])
            ellipse2 = tuple(draw_ellipse_requirements[body_idx_2])

            if not contains(ellipse1, ellipse2):
                continue

            # is_overlap = compute_overlap(ellipse_boxes[body_idx_1], ellipse_boxes[body_idx_2], use_ell_box=False)
            # if is_overlap:
            is_skeletal_overlapped[body_idx_1] = 1
            is_skeletal_overlapped[body_idx_2] = 1
            break

    colors = [(0, 255, 0), (0, 0, 255)]  # 0 = negtive = green, 1 = positive = red
    for i in range(len(draw_ellipse_requirements)):
        color = colors[int(is_skeletal_overlapped[i])]
        im = cv2.ellipse(im, (draw_ellipse_requirements[i][0], draw_ellipse_requirements[i][1]),
                         (int(draw_ellipse_requirements[i][2]), draw_ellipse_requirements[i][3]), 0, 0, 360,
                         color,2)
        im = cv2.circle(im,(draw_ellipse_requirements[i][0], draw_ellipse_requirements[i][1]),5,color,-1)

    return is_skeletal_overlapped,im

if __name__=='__main__':

    model_path = "yolov5s.pt" # it automatically downloads yolov5s model to given path
    device = "cuda" # or "cpu"  cuda

    # capture = cv2.VideoCapture('TownCentre.avi')
    capture = cv2.VideoCapture('x.avi')

    yolov5 = YOLOv5(model_path, device)

    horizontal_ratio, vertical_ratio=0.6,0.5
    body_height = 1.7

    scale=1.0
    H_truth = np.array([[0.000161618271024591, -0.000517859162483943, 0.5583877529554223],
                        [-0.00033418918205856775, -0.0002735434211166214, 0.6920511725385885],
                        [1.113222825743109e-06, 4.4403770399997656e-05, 0.018654311930500225]])

    frame_i=0

    while True:
        _, im = capture.read()
        if im is None:
            break
        frame_i=frame_i+1

        start_time = time.perf_counter()

        results = yolov5.predict(im,size=320)
        bboxes = []
        for result in results.xyxy[0]:
            result= result.cpu().numpy().tolist()
            x1,y1,x2,y2,score,cls, = result[:]
            x1,y1,x2,y2 = int(x1),int(y1),int(x2),int(y2)
            if cls!=0:
                continue
            bboxes.append((x1, y1, x2, y2, 'person', score))
        # pdb.set_trace()
        if len(bboxes) > 0:
            draw_ellipse_requirements,ellipse_boxes,center_world_coordinates = get_ells_from_person_bboxes(im,bboxes,horizontal_ratio, vertical_ratio)

            #绘制ell, 预测overlap_pre
            is_skeletal_overlapped = overlap_infer(im, draw_ellipse_requirements, ellipse_boxes)
            # calculate fps

            elapsed = time.perf_counter() - start_time
            cv2.putText(im, "fps= {:.3f}".format(1/elapsed)
                    # "frame:{}-presise={}/({}+{})={:.3f},recall={}/({}+{})={:.3f},F1={:.3f}".format(
                    # str(frame_i),TPall,TPall,FPall,precise,TPall,TPall,FNall,recall,F1)
                    , (20, 80), 0, 1, (0, 255,0), 3)

        cv2.imshow('d',im);cv2.imwrite('d.jpg',im)
        cv2.waitKey()



